'''
Author: SlytherinGe
LastEditTime: 2021-03-15 19:15:04
'''
import numpy as np
import os
import time
import sys
import xml.etree.ElementTree as ET
import math

if __name__ == '__main__':
    import voc_label_utility as util
else:
    import dataset_utilty.voc_label_utility as util


def make_small_bboxes_bigger(in_root, out_root, area_of_small, biggger_box_size):
    anno_files, anno_list = util.load_voc_annotations(in_root)
    for i, anno in enumerate(anno_list):
        im_h = int(anno['size']['height'])
        im_w = int(anno['size']['width'])
        for k, bndbox in enumerate(anno['object']):
            l = math.sqrt(area_of_small)
            bbox = (int(bndbox['bndbox']['xmin']), int(bndbox['bndbox']['ymin']), int(bndbox['bndbox']['xmax']), int(bndbox['bndbox']['ymax'])) 
            if (bbox[2]-bbox[0]) < l and (bbox[3] - bbox[1]) < l:
                temp = util.lefttop_rightdown_2_center_width(bbox)
                bigger_box = (temp[0] - biggger_box_size/2, temp[1] - biggger_box_size/2, temp[0] + biggger_box_size/2, temp[1] + biggger_box_size/2)
                bigger_box = util.clip_bbox(bigger_box, (im_w, im_h))
                anno_list[i]['object'][k]['bndbox']['xmin'] = str(bigger_box[0])
                anno_list[i]['object'][k]['bndbox']['ymin'] = str(bigger_box[1])
                anno_list[i]['object'][k]['bndbox']['xmax'] = str(bigger_box[2])
                anno_list[i]['object'][k]['bndbox']['ymax'] = str(bigger_box[3])
                anno_list[i]['object'][k]['name'] = 'small'
    util.save_anno_list(anno_files, anno_list, out_root, 0)

if __name__ == '__main__':
    ANNO_ROOT = '/media/gejunyao/Disk1/Datasets/SARship-Total/PASCAL_VOC/VOC2012/Annotations_ships'
    OUT_ROOT = '/media/gejunyao/Disk1/Datasets/SARship-Total/PASCAL_VOC/VOC2012/Annotations'
    make_small_bboxes_bigger(ANNO_ROOT, OUT_ROOT, 1250, 50)